import openpyxl
from flask import Flask, request, jsonify
import scipy.stats as stats
from database import pool
import traceback
import re

app = Flask(__name__)

heat_change_coefficient = 0  # 专业热度变化调整系数，默认为0  >0热度升高   <0 热度降低


def calculate_admission_probability(enrollment_plan, candidate_rank,
                                    historical_min_rank_list, historical_enrollment_list,
                                    heat_change=0):
    """
    计算考生被专业录取的概率

    :param enrollment_plan: 目标高校专业A在上海的招生计划人数
    :param candidate_rank: 考生个人的总排名
    :param historical_min_rank_list: 历年录取最低位次数据列表
    :param historical_enrollment_list: 历年录取人数数据列表
    :param heat_change: 专业热度变化调整系数，默认为0
    :return: 考生被专业录取的概率
    """
    # 计算历年录取最低位次的统计量
    # 计算平均最低位次
    average_min_rank = sum(historical_min_rank_list) / len(historical_min_rank_list)
    # 计算标准差
    variance = sum((x - average_min_rank) ** 2 for x in historical_min_rank_list) / len(historical_min_rank_list)
    standard_deviation = variance ** 0.5

    # 构建概率计算模型（基于正态分布假设）
    # 计算标准分数（Z - score）
    if standard_deviation == 0:
        # 处理标准差为 0 的情况
        distance = abs(candidate_rank - average_min_rank)
        if candidate_rank < average_min_rank:
            # 当考生排名小于平均最低位次时，距离越大z_score越小
            z_score = -distance / (max(historical_min_rank_list) - min(historical_min_rank_list) + 1)
            # print('candidate_rank < average_min_rank', z_score)
        elif candidate_rank > average_min_rank:
            # 当考生排名大于平均最低位次时，距离越大z_score越大
            z_score = distance / (max(historical_min_rank_list) - min(historical_min_rank_list) + 1)
            # print('candidate_rank > average_min_rank', z_score)
        else:
            z_score = 0
            # print('candidate_rank == average_min_rank', z_score)
    else:
        z_score = (candidate_rank - average_min_rank) / standard_deviation
        # print('z_score:', z_score)

    # 计算录取概率
    base_probability = stats.norm.sf(z_score)

    # 考虑招生计划调整概率
    # 计算历年平均录取人数
    average_historical_enrollment = sum(historical_enrollment_list) / len(historical_enrollment_list)
    # 调整因子
    adjustment_factor = enrollment_plan / average_historical_enrollment
    # 调整后的概率
    adjusted_admission_probability = base_probability * adjustment_factor
    # 确保概率在[0, 1]范围内
    adjusted_admission_probability = max(0, min(adjusted_admission_probability, 1))

    # 考虑专业热度变化因素
    # 调整后的概率
    final_admission_probability = adjusted_admission_probability * (1 + heat_change)
    # 确保最终概率在[0, 1]范围内
    final_admission_probability = max(0, min(final_admission_probability, 1))

    return final_admission_probability



def export_result_to_excel(name, score, rank, subjects, data):
    # 创建一个空的DataFrame，用于存储结果
    import pandas as pd
    from datetime import datetime
    from openpyxl import Workbook
    from openpyxl.styles import Alignment
    from openpyxl.utils.dataframe import dataframe_to_rows

    # 第一行为文件表头：专业概率表 姓名：name ,总分：score,排名：rank, subjects
    header = f"专业概率表 姓名：{name}, 总分：{score}, 排名：{rank}, 科目：{subjects}"

    # 第二行为数据表头：'学校代码', '学校', '专业组', '专业代码', '专业', '招生人数', '录取概率'
    columns = ['学校代码', '学校', '专业组', '专业代码', '专业', '招生人数', '录取概率']

    # 第三行开始为数据
    df = pd.DataFrame(data, columns=columns)

    # 将录取概率转换为百分比形式
    df['录取概率'] = df['录取概率'].apply(lambda x: f"{(x * 100):.2f}%")

    # 文件名带上日期时间
    filename = f'results_{datetime.now().strftime("%Y%m%d%H%M%S")}.xlsx'

    # 使用openpyxl创建一个新的工作簿
    wb = Workbook()
    ws = wb.active

    # 写入表头
    ws.append([header])

    # 写入数据表头
    ws.append(columns)

    # 将DataFrame写入工作表，从第三行开始
    for r in dataframe_to_rows(df, index=False, header=False):
        ws.append(r)

    # 自动调整列宽
    for column_cells in ws.columns:
        # 获取列号
        column_letter = column_cells[0].column_letter
        # 找到非合并单元格的最大宽度，排除第一行的数据
        length = max(len(str(cell.value)) for cell in column_cells if cell.row != 1)
        ws.column_dimensions[column_letter].width = length + 5

    # 合并第一行的7个单元格
    ws.merge_cells('A1:G1')

    # 居中对齐第一行的单元格
    cell = ws['A1']
    cell.alignment = Alignment(horizontal='center', vertical='center')

    # 保存工作簿到文件
    wb.save(f'file/{filename}')
    print(f'导出成功，文件名: {filename}')


@app.route('/api/admission_probability', methods=['POST'])
def admission_probability():
    # 获取请求中的考生信息
    data = request.get_json()
    year = data.get('year')
    source_place = data.get('source_place')
    exam_subjects = data.get('exam_subjects')
    total_score = data.get('total_score')
    rank = data.get('rank')
    name = data.get('name')

    print('请求中的考生信息：', data)

    db = pool.get_connection()

    cursor = db.cursor(dictionary=True)

    # 查询招生计划
    query = """
    SELECT id,school_code,school,province,city,soft_rank,is_985,is_211,is_double_first_class,
    batch,major_group,exam_require,exam_require_regex,category,
    first_subject,major,major_code,enroll_num
    FROM t_enroll_plan
    WHERE year = %s AND source_place = %s  AND  school_code IS NOT NULL 
          AND  major_code IS NOT NULL AND  enroll_num > 0
    ORDER BY is_985 desc, is_211 desc , soft_rank
    """
    cursor.execute(query, (year, source_place))
    enroll_plans = cursor.fetchall()

    huandred_percent_num = 0
    # 计算录取概率
    results = []
    for enroll_plan in enroll_plans:
        # 100%录取人数>=20，后续不要算了
        if huandred_percent_num >= 100:
            break

        school_code = enroll_plan['school_code']
        school = enroll_plan['school']
        major_group = enroll_plan['major_group']
        major_code = enroll_plan['major_code']
        major = enroll_plan['major']
        enroll_num = enroll_plan['enroll_num']
        exam_require = enroll_plan['exam_require']
        exam_require_regex = enroll_plan['exam_require_regex']

        # print('正在查询', source_place, school, school_code, major_group, major_code, major, '的录取概率')

        if exam_require_regex is not None:
            match = re.match(exam_require_regex, exam_subjects)
            if not match:
                # print('考试科目要求不能匹配: ', source_place, school, school_code, major_group, major_code, major, exam_require,
                #       exam_require_regex)
                continue
        # 查询最近5年的相关数据
        query = """
        SELECT 
            sl.min_rank,
            sl.num
        FROM 
            t_score_line sl
        WHERE 
            sl.school_code = %s AND sl.major_code = %s AND sl.source_place = %s   AND sl.min_rank > 0  AND sl.num > 0
        ORDER BY 
            sl.year 
        """
        cursor.execute(query, (school_code, major_code, source_place))
        stats = cursor.fetchall()

        if stats:
            wcss = [stat['min_rank'] for stat in stats]
            sums = [stat['num'] for stat in stats]
            # print(f'历年该专业录取最低分对应的位次wcss：{wcss}')
            # print(f'历年该专业录取人数sums：{sums}')
            # print(f'今年录取人数：{enroll_num}')
        else:
            print(source_place, school, school_code, major_group, major_code, major, '历史数据无，无法估计概率')
            continue

        # if wcss[0] >= rank:
            # print(f'该考生的录取位次小于等于录取最低位次: {wcss[0]} >= {rank}')

        try:
            probability = calculate_admission_probability(enroll_num, rank, wcss, sums, heat_change_coefficient)

            if probability >= 0.99:  # 录取概率100%
                huandred_percent_num += 1
            elif probability < 0.01:  # 录取概率小于1%，跳过该专业
                continue

        except Exception as e:
            print(f"An error occurred: {e}")
            traceback_str = traceback.format_exc()
            print(traceback_str)
            probability = 0
            print("位次:", rank)
            print("专业录取人数:", enroll_num)

        # print(source_place, school, school_code, major_group, major_code, major, probability)

        results.append([school_code, school, major_group, major_code, major, enroll_num, probability])

    # 关闭数据库连接
    cursor.close()
    db.close()

    # 按照录取概率升序排序
    results.sort(key=lambda x: x[6], reverse=False)

    # 将results导出到excel文件中,其中概率按照百分比输出
    if results:
        export_result_to_excel(name, total_score, rank, exam_subjects, results)

    json_response = []
    for school_code, school, major_group, major_code, major, enroll_num, probability in results:
        json_response.append({
            '学校代码': school_code,
            '学校': school,
            '专业组': major_group,
            '专业代码': major_code,
            '专业': major,
            '录取人数': enroll_num,
            '录取概率': f"{probability:.2f}"
        })

    return jsonify(json_response)


@app.route('/')
def home():
    return "Hello, Flask!"


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)
